import numpy as np
import cv2
import pickle
import matplotlib.pyplot as plt
import logging
from create_clf import create_features

# 配置日志记录
logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s - %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

def main():
    """
    主程序入口，协调砂岩图像分割流程
    """
    logger.info("开始砂岩图像分割流程...")
    
    try:
        # 加载模型
        with open('sandstone_clf.pkl', 'rb') as f:
            clf = pickle.load(f)
        logger.info("成功从硬盘加载随机森林模型clf")
        
        # 读取图像
        image = cv2.imread('Sandstone_2.tif', cv2.IMREAD_GRAYSCALE)
        true_segment = cv2.imread('Sandstone_2_segment.tif', cv2.IMREAD_GRAYSCALE)
        
        if image is None or true_segment is None:
            raise FileNotFoundError("未找到图像文件")
        
        logger.info(f"图像形状: {image.shape}")
        
        # 创建特征
        X = create_features(image)
        logger.info("完成从新砂岩图像中获取特征矩阵X")
        
        # 预测
        predictions = clf.predict(X)
        predicted_image = predictions.reshape(image.shape)
        logger.info("完成图像分割")
        
        # 计算准确率
        accuracy = np.mean(predicted_image.ravel() == true_segment.ravel())
        logger.info(f"分割准确率: {accuracy:.2%}")
        
        # 显示结果
        plt.figure(figsize=(15, 5))
        
        plt.subplot(131)
        plt.imshow(image, cmap='gray')
        plt.title('Original')
        plt.axis('off')
        
        plt.subplot(132)
        plt.imshow(true_segment, cmap='gray')
        plt.title('True')
        plt.axis('off')
        
        plt.subplot(133)
        plt.imshow(predicted_image, cmap='gray')
        plt.title(f'Predict(Accuracy: {accuracy:.2%})')
        plt.axis('off')
        
        plt.tight_layout()
        plt.show()
        logger.info("分割结果可视化完成")
        
    except Exception as e:
        logger.error(f"处理过程中发生错误: {str(e)}")

if __name__ == "__main__":
    main()